package ratelimit.single;

import com.google.common.util.concurrent.RateLimiter;

import java.util.WeakHashMap;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantReadWriteLock;


/**
 * @Auther: Fazhan.Qiao
 * @Date: 2019/5/27 10:04
 * @Description: 单机应用平滑突发限流策略（令牌桶算法）
 * 可以针对某个单机应用进行请求总量的限流，用在Filter中；
 * 也可以针对某个方法进行限流，如用在Controller、service中；
 * demo()方法演示了如何使用。mian方法是一些测试
 *
 * 不同业务的限流名称不要相同，名称相同会使用同一个限流实例；
 */
public class TokenLimitFacory {
    private static WeakHashMap<String, RateLimiter> limiterMap = new WeakHashMap<>();
    private static ReentrantReadWriteLock lock = new ReentrantReadWriteLock();

    public static void acquire(String keyName, int limit) {//限流后阻塞（排队等待）
        RateLimiter rateLimiter=null;
        try {
            lock.readLock().lock();
            if (limiterMap.containsKey(keyName)) {
                rateLimiter = limiterMap.get(keyName);
            }
        } catch (Exception e){
            e.printStackTrace();
        }finally {
            lock.readLock().unlock();
        }

        if (rateLimiter == null) {
            try {
                lock.writeLock().lock();
                if (limiterMap.containsKey(keyName)) {//双重检查
                    rateLimiter = limiterMap.get(keyName);
                }
                if (rateLimiter == null) {
                    rateLimiter = RateLimiter.create(limit);
                    limiterMap.put(keyName, rateLimiter);
                }
            }catch (Exception e){
                e.printStackTrace();
            }finally {
                lock.writeLock().unlock();
            }
        }
        rateLimiter.acquire();
    }

    public static boolean tryAcquire(String keyName, int limit) {//限流后拒绝
        RateLimiter rateLimiter=null;
        try {
            lock.readLock().lock();
            if (limiterMap.containsKey(keyName)) {
                rateLimiter = limiterMap.get(keyName);
            }
        } catch (Exception e){
            e.printStackTrace();
        }finally {
            lock.readLock().unlock();
        }

        if (rateLimiter == null) {
            try {
                lock.writeLock().lock();
                if (limiterMap.containsKey(keyName)) {//双重检查
                    rateLimiter = limiterMap.get(keyName);
                }
                if (rateLimiter == null) {
                    rateLimiter = RateLimiter.create(limit);
                    limiterMap.put(keyName, rateLimiter);
                }
            }catch (Exception e){
                e.printStackTrace();
            }finally {
                lock.writeLock().unlock();
            }
        }
        return rateLimiter.tryAcquire(1000 * 1000 * 1000 - 1000 * 1000 * 1000 / limit, TimeUnit.NANOSECONDS);
    }


/*    public static void demo1() {
        //限流每秒1000
        TokenLimitFacory.acquire("user.query",1000);//获取不到token时，会一直阻塞；请求排队等待；
         // 处理核心逻辑

    }*/



/*    public static void demo2() {
          //限流每秒1000
       boolean nolimit= TokenLimitFacory.tryAcquire("user.query",1000);//获取不到token时，会直接返回false；
        if(nolimit){//获取到token了
           // 处理核心逻辑
        }else{
            //没有获取到token  限流了，直接丢弃请求，抛出异常
            throw new RuntimeException("当前请求过多，请稍后再试")
        }


    }*/


    public static void main(String[] args) {

        ExecutorService threadPoolExecutor = new ThreadPoolExecutor(5, 200, 30L, TimeUnit.MINUTES, new ArrayBlockingQueue<Runnable>(50));

//        threadPoolExecutor.execute(new Runnable() {
//            @Override
//            public void run() {
//                for (int i = 0; i < 50; i++) {
//                    System.out.println("1111");
//                }
//            }
//        });

//        for (int i = 0; i < 200; i++) {
//            threadPoolExecutor.execute(new Runnable() {
//                @Override
//                public void run() {
//                    //临牌桶中1秒中放入12个token，并发200个请求，只有12个请求能通过，其他限流后直接拒绝
//                    boolean noLimit = TokenLimitFacory.tryAcquire("print2", 12);
//                    if (noLimit) {
//                        System.out.println("22222222   :");
//                    }
//                }
//            });
//        }


        for (int i = 0; i < 100; i++) {
            threadPoolExecutor.execute(new Runnable() {
                @Override
                public void run() {
                    //临牌桶中1秒中放入12个token，并发100个请求，这100个请求会每秒钟处理12个，其他的请求排队等待；10s处理完所有请求
                    TokenLimitFacory.acquire("print3", 12);
                    System.out.println("33333333333333   :");
                }
            });
        }


    }
}
